"""
框架2：粮食（直接预测）
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
import sys
import os

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ..data_loader import DataLoader
from ..ilr_transformation import ILRTransformation
from ..forecasting_models import ForecastingModels
from ..evaluation import Evaluation


class Framework2:
    """框架2：粮食（直接预测）"""
    
    def __init__(self, data: Dict, config: Dict = None):
        """
        初始化框架2
        
        Args:
            data: 处理后的数据字典
            config: 配置字典
        """
        self.data = data
        self.config = config or {}
        self.forecasting_models = ForecastingModels(config)
        self.evaluation = Evaluation()
        
        # 提取数据
        self._extract_data()
        
    def _extract_data(self):
        """提取和准备数据"""
        # 获取第一个工作表的数据
        sheet_name = list(self.data.keys())[0]
        sheet_data = self.data[sheet_name]
        
        # 提取比例数据
        self.proportions_data = sheet_data['proportions']
        
        # 提取各省份在全国的占比数据
        self.province_national_prop = self.proportions_data['province_national']
        
        # 提取全国总量数据
        self.national_data = sheet_data['aggregated']['national']
        
    def prepare_compositional_data(self) -> pd.DataFrame:
        """
        准备成分数据
        
        Returns:
            成分数据矩阵
        """
        # 准备各省份在全国的占比数据
        province_wide = self.province_national_prop.pivot_table(
            index='Year',
            columns='Province',
            values='National_Proportion',
            fill_value=0
        )
        
        # 处理零值和负值问题
        epsilon = 1e-10  # 小的正数替代零值
        
        # 将零值和负值替换为小正数
        province_wide = province_wide.where(province_wide > 0, epsilon)
        
        # 确保每行和为1（重新归一化）
        row_sums = province_wide.sum(axis=1)
        province_wide = province_wide.div(row_sums, axis=0)
        
        # 转换为矩阵形式
        mat = province_wide.values
        
        try:
            # 对成分数据进行ILR变换
            ilr_transformer = ILRTransformation()
            ilr_transformer.fit(mat)
            ilr_coords = ilr_transformer.transform(mat)
            
            # 创建最终数据框
            ilr_df = pd.DataFrame(ilr_coords, columns=[f'ILR_{i+1}' for i in range(ilr_coords.shape[1])])
            ilr_df['Year'] = province_wide.index
            
            # 合并全国总量数据和ILR坐标
            national_df = self.national_data[['Year', 'Grain']].copy()
            final_data = national_df.merge(ilr_df, on='Year')
            
            return final_data
            
        except Exception as e:
            print(f"ILR变换失败: {e}")
            # 如果ILR变换失败，返回简化的数据
            simplified_data = self.national_data[['Year', 'Grain']].copy()
            # 添加主要省份的比例作为特征
            main_provinces = province_wide.columns[:5]  # 取前5个省份
            for province in main_provinces:
                simplified_data[f'{province}_prop'] = province_wide[province].values
            
            return simplified_data
    
    def forecast(self, horizon: int = 5) -> pd.DataFrame:
        """
        执行简化的预测
        
        Args:
            horizon: 预测步长
            
        Returns:
            预测结果DataFrame
        """
        try:
            from statsmodels.tsa.arima.model import ARIMA
            import numpy as np
            
            print("Framework2: 开始MinT协调预测...")
            
            # 使用简化的方法：直接对省份数据进行预测，然后协调
            province_data = self.province_national_prop.pivot_table(
                index='Year',
                columns='Province', 
                values='National_Proportion',
                fill_value=0
            )
            
            # 处理数据质量问题
            epsilon = 1e-10
            province_data = province_data.where(province_data > 0, epsilon)
            
            # 重新归一化确保比例和为1
            row_sums = province_data.sum(axis=1)
            province_data = province_data.div(row_sums, axis=0)
            
            print(f"  省份数据形状: {province_data.shape}")
            
            # 对每个省份的比例进行ARIMA预测
            forecasts = {}
            
            for province in province_data.columns:
                series = province_data[province].dropna()
                
                if len(series) < 10:
                    print(f"  跳过 {province}：数据点不足")
                    continue
                
                try:
                    # 使用简单的ARIMA模型预测比例
                    model = ARIMA(series, order=(1, 1, 1))
                    fitted_model = model.fit()
                    forecast_result = fitted_model.forecast(steps=horizon)
                    
                    # 确保预测的比例为正数
                    forecast_result = np.maximum(forecast_result, epsilon)
                    forecasts[province] = forecast_result.values
                    
                except Exception as e:
                    print(f"  预测 {province} 失败: {str(e)}")
                    # 使用最后几年的平均值作为预测
                    recent_mean = series.tail(3).mean()
                    forecasts[province] = np.full(horizon, max(recent_mean, epsilon))
            
            if not forecasts:
                print("  没有成功预测任何省份")
                return None
            
            # 构建预测结果DataFrame
            forecast_df = pd.DataFrame(forecasts)
            
            # 重新归一化确保每行比例和为1
            row_sums = forecast_df.sum(axis=1)
            forecast_df = forecast_df.div(row_sums, axis=0)
            
            # 预测全国总量
            national_data = self.national_data.set_index('Year')['Grain']
            
            try:
                # 预测全国总量
                national_model = ARIMA(national_data, order=(1, 1, 1))
                national_fitted = national_model.fit()
                national_forecast = national_fitted.forecast(steps=horizon)
                
                print(f"  全国总量预测: {national_forecast.values}")
                
            except Exception as e:
                print(f"  全国总量预测失败，使用趋势外推: {str(e)}")
                # 使用简单的趋势外推
                recent_growth = national_data.tail(5).pct_change().mean()
                last_value = national_data.iloc[-1]
                national_forecast = [last_value * (1 + recent_growth) ** (i+1) for i in range(horizon)]
                national_forecast = pd.Series(national_forecast)
            
            # 将比例转换为绝对值
            result_df = forecast_df.copy()
            for i in range(horizon):
                result_df.iloc[i] = result_df.iloc[i] * national_forecast.iloc[i]
            
            # 设置正确的时间索引
            last_year = province_data.index.max()
            future_years = range(last_year + 1, last_year + horizon + 1)
            result_df.index = future_years
            
            print(f"  Framework2预测完成，结果形状: {result_df.shape}")
            return result_df
            
        except Exception as e:
            print(f"Framework2预测失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def _reconstruct_forecasts(self, forecast_results: pd.DataFrame, horizon: int) -> Dict:
        """
        重建预测结果
        
        Args:
            forecast_results: 预测结果数据框
            horizon: 预测步长
            
        Returns:
            重建后的预测结果字典
        """
        # 分离全国总量预测和ILR预测
        national_forecast = forecast_results['Grain'].values
        ilr_forecasts = forecast_results[[col for col in forecast_results.columns if 'ILR' in col]]
        
        # 对ILR预测进行逆变换
        ilr_transformer = ILRTransformation()
        # 使用训练数据拟合变换器
        train_compositional = self.prepare_compositional_data()
        train_mat = train_compositional.drop(['Year', 'Grain'], axis=1).values
        ilr_transformer.fit(train_mat)
        
        # 逆变换
        reconstructed_proportions = ilr_transformer.inverse_transform(ilr_forecasts.values)
        
        # 重建各省份预测
        province_forecasts = {}
        provinces = self.province_national_prop['Province'].unique()
        
        for i, province in enumerate(provinces):
            if i < reconstructed_proportions.shape[1]:
                province_forecasts[province] = national_forecast * reconstructed_proportions[:, i]
        
        # 重建区域预测
        regional_forecasts = self._aggregate_to_regions(province_forecasts)
        
        return {
            'national': national_forecast,
            'regional': regional_forecasts,
            'province': province_forecasts
        }
    
    def _aggregate_to_regions(self, province_forecasts: Dict) -> Dict:
        """
        聚合省份预测到区域预测
        
        Args:
            province_forecasts: 省份预测字典
            
        Returns:
            区域预测字典
        """
        regional_forecasts = {}
        
        # 定义区域-省份映射
        region_provinces = {
            '华北地区': ['北京市', '河北+天津', '山西省', '内蒙古自治区'],
            '东北地区': ['辽宁省', '吉林省', '黑龙江省'],
            '华东地区': ['上海市', '江苏省', '浙江省', '安徽省', '福建省', '江西省', '山东省'],
            '华中地区': ['河南省', '湖北省', '湖南省'],
            '华南地区': ['广东+海南', '广西壮族自治区'],
            '西南地区': ['四川+重庆', '贵州省', '云南省', '西藏自治区'],
            '西北地区': ['陕西省', '甘肃省', '青海省', '宁夏回族自治区', '新疆维吾尔自治区']
        }
        
        for region, provinces in region_provinces.items():
            region_total = np.zeros(len(next(iter(province_forecasts.values()))))
            for province in provinces:
                if province in province_forecasts:
                    region_total += province_forecasts[province]
            regional_forecasts[region] = region_total
        
        return regional_forecasts
    
    def run_complete_analysis(self, horizons: List[int] = None) -> Dict:
        """
        运行完整的分析流程
        
        Args:
            horizons: 预测步长列表
            
        Returns:
            完整分析结果
        """
        if horizons is None:
            horizons = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 15, 20]
        
        all_results = {}
        
        for horizon in horizons:
            print(f"正在处理预测步长 h={horizon}")
            
            # 执行预测
            forecast_results = self.forecast(horizon=horizon)
            
            all_results[f"h={horizon}"] = {
                'forecast': forecast_results
            }
        
        return all_results 